import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from FlagEmbedding import BGEM3FlagModel
import numpy as np
import faiss
from sklearn.metrics.pairwise import cosine_similarity
import torch

# 加载BGE-M3模型
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)

# 示例文本数据
corpus = [
    "自然语言处理是人工智能的一个重要领域",
    "大型语言模型在许多NLP任务中表现出色",
    "计算机视觉处理图像和视频数据",
    "推荐系统基于用户行为预测偏好",
    "信息检索关注如何高效获取相关信息"
]

# 查询文本
query = "什么技术用于文本分析？"

# 生成文本向量（密集检索）
def get_dense_embeddings(texts):
    embeddings = model.encode(texts, batch_size=128, max_length=8192)['dense_vecs']
    # 确保向量是numpy数组且为float32类型（FAISS要求）
    return np.array(embeddings).astype('float32')

# 生成稀疏向量（稀疏检索）
def get_sparse_embeddings(texts):
    output = model.encode(texts, return_dense=True, return_sparse=True, return_colbert_vecs=False)
    return output['lexical_weights']

# 使用FAISS构建索引并执行密集检索
class FaissIndex:
    def __init__(self, dim=768):
        # 创建FAISS索引（使用内积相似度，等价于归一化后的余弦相似度）
        self.index = faiss.IndexFlatIP(dim)
        self.documents = []
        
    def add_vectors(self, vectors, documents):
        # 向量归一化（对于余弦相似度计算很重要）
        faiss.normalize_L2(vectors)
        self.index.add(vectors)
        self.documents = documents
        
    def search(self, query_vector, k=5):
        # 查询向量归一化
        query_vector = query_vector.copy().astype('float32')
        faiss.normalize_L2(query_vector)
        
        # 执行搜索
        distances, indices = self.index.search(query_vector, k)
        
        # 返回结果
        results = []
        for i, idx in enumerate(indices[0]):
            if idx != -1:  # -1表示没有找到
                results.append((self.documents[idx], distances[0][i]))
        return results

# 密集检索示例
def dense_retrieval(query, corpus):
    # 生成查询和语料库的密集向量
    query_embedding = get_dense_embeddings([query])
    corpus_embeddings = get_dense_embeddings(corpus)
    
    # 初始化FAISS索引
    index = FaissIndex(dim=corpus_embeddings.shape[1])
    index.add_vectors(corpus_embeddings, corpus)
    
    # 执行搜索
    results = index.search(query_embedding, k=len(corpus))
    
    return results

# 稀疏检索示例
def sparse_retrieval(query, corpus):
    # 生成查询和语料库的稀疏向量
    query_sparse = get_sparse_embeddings([query])[0]
    corpus_sparse = get_sparse_embeddings(corpus)
    
    # 计算词法匹配分数
    scores = []
    for doc_sparse in corpus_sparse:
        score = model.compute_lexical_matching_score(query_sparse, doc_sparse)
        scores.append(score)
    
    # 按分数排序
    results = sorted(zip(corpus, scores), key=lambda x: x[1], reverse=True)
    
    return results

# 混合检索示例（结合密集和稀疏检索）
def hybrid_retrieval(query, corpus, alpha=0.7):
    # 密集检索结果
    dense_results = dense_retrieval(query, corpus)
    dense_scores = {doc: score for doc, score in dense_results}
    
    # 稀疏检索结果
    sparse_results = sparse_retrieval(query, corpus)
    sparse_scores = {doc: score for doc, score in sparse_results}
    
    # 混合分数
    hybrid_scores = []
    for doc in corpus:
        # 将稀疏分数归一化到[0,1]范围
        norm_sparse_score = (sparse_scores[doc] - min(sparse_scores.values())) / \
                           (max(sparse_scores.values()) - min(sparse_scores.values()) + 1e-8)
        hybrid_score = alpha * dense_scores[doc] + (1 - alpha) * norm_sparse_score
        hybrid_scores.append((doc, hybrid_score))
    
    # 按混合分数排序
    results = sorted(hybrid_scores, key=lambda x: x[1], reverse=True)
    
    return results

# 执行检索
print("===== 密集检索结果 =====")
for doc, score in dense_retrieval(query, corpus):
    print(f"相似度: {score:.4f} | 文档: {doc}")

print("\n===== 稀疏检索结果 =====")
for doc, score in sparse_retrieval(query, corpus):
    print(f"相似度: {score:.4f} | 文档: {doc}")

print("\n===== 混合检索结果 =====")
for doc, score in hybrid_retrieval(query, corpus):
    print(f"相似度: {score:.4f} | 文档: {doc}")